import os
import logging
import requests
import base64
import json
import tempfile
import numpy as np
from typing import Dict, Any, Optional
from datetime import datetime
from concurrent.futures import ThreadPoolExecutor

# 添加实际渲染所需的导入
try:
    import trimesh
    import pyrender
    from pyrender import PerspectiveCamera, DirectionalLight, OffscreenRenderer
    import matplotlib.pyplot as plt
    PYRENDER_AVAILABLE = True
except ImportError as e:
    logging.warning(f"无法导入pyrender或相关库: {str(e)}")
    logging.warning("将使用模拟渲染模式")
    PYRENDER_AVAILABLE = False

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
logger = logging.getLogger(__name__)

# 使用实际pyrender进行3D模型切图
class ModelRenderer:
    def __init__(self):
        logger.info(f"初始化3D模型渲染器，Pyrender可用状态: {PYRENDER_AVAILABLE}")
        
        if PYRENDER_AVAILABLE:
            # 初始化pyrender环境
            self.renderer = OffscreenRenderer(viewport_width=640, viewport_height=480)
            self.screenshot_dir = os.path.join(tempfile.gettempdir(), "3d_model_screenshots")
            os.makedirs(self.screenshot_dir, exist_ok=True)
        else:
            # 当pyrender不可用时使用模拟实现
            self.mock_screenshot_url = "https://example.com/mock_screenshot.jpg"
            self.screenshot_dir = None
        
        # 无论是否可用，都初始化mock_screenshot_url作为后备
        self.mock_screenshot_url = "https://example.com/mock_screenshot.jpg"
    
    def render_model(self, model_url: str, output_dir: Optional[str] = None) -> Dict[str, Any]:
        """
        渲染3D模型并生成多角度截图
        
        Args:
            model_url: 3D模型URL
            output_dir: 输出目录
            
        Returns:
            包含截图信息的字典
        """
        try:
            logger.info(f"开始渲染3D模型: {model_url}")
            
            if not PYRENDER_AVAILABLE:
                # 返回模拟的截图结果
                return self._mock_render(model_url)
            
            # 创建输出目录
            if output_dir:
                self.screenshot_dir = output_dir
                os.makedirs(self.screenshot_dir, exist_ok=True)
            
            # 下载模型文件
            model_path = self._download_model(model_url)
            if not model_path:
                return self._mock_render(model_url, error="无法下载模型文件")
            
            # 加载模型
            mesh = self._load_model(model_path)
            if not mesh:
                return self._mock_render(model_url, error="无法加载模型")
            
            # 从多个角度渲染模型
            screenshots = self._render_multiple_views(mesh)
            
            logger.info("3D模型渲染完成")
            return screenshots
            
        except Exception as e:
            logger.error(f"3D模型渲染失败: {str(e)}")
            # 发生异常时返回备用数据
            return self._mock_render(model_url, error=str(e))
    
    def _download_model(self, model_url: str) -> Optional[str]:
        """
        下载3D模型文件
        """
        try:
            # 从URL中提取文件名，忽略查询参数
            url_path = model_url.split('?')[0]
            file_extension = os.path.splitext(url_path)[1] or '.obj'
            # 确保扩展名是有效的3D模型格式
            if file_extension not in ['.glb', '.gltf', '.obj', '.fbx', '.stl']:
                file_extension = '.glb'
            
            # 创建临时文件
            temp_file = tempfile.NamedTemporaryFile(suffix=file_extension, delete=False)
            temp_file_path = temp_file.name
            temp_file.close()
            
            logger.info(f"正在下载模型到临时文件: {temp_file_path}")
            
            # 下载模型
            response = requests.get(model_url, stream=True, timeout=60)
            response.raise_for_status()
            
            with open(temp_file_path, 'wb') as f:
                for chunk in response.iter_content(chunk_size=8192):
                    f.write(chunk)
            
            logger.info(f"模型下载完成，文件大小: {os.path.getsize(temp_file_path)} 字节")
            return temp_file_path
            
        except Exception as e:
            logger.error(f"下载模型失败: {str(e)}")
            return None
    
    def _load_model(self, model_path: str) -> Optional[trimesh.Trimesh]:
        """
        加载3D模型
        """
        try:
            # 使用trimesh加载模型
            mesh = trimesh.load(model_path)
            
            # 如果是场景，提取第一个网格
            if isinstance(mesh, trimesh.Scene):
                geometries = list(mesh.geometry.values())
                if geometries:
                    mesh = geometries[0]
                else:
                    return None
            
            return mesh
            
        except Exception as e:
            logger.error(f"加载模型失败: {str(e)}")
            return None
    
    def _render_multiple_views(self, mesh: trimesh.Trimesh) -> Dict[str, Any]:
        """
        从多个角度渲染模型
        """
        try:
            screenshots = {}
            render_status = "success"
            
            # 定义不同视角
            view_angles = {
                "front": np.array([0, 0, np.pi]),
                "side": np.array([0, 0, np.pi/2]),
                "top": np.array([np.pi/2, 0, np.pi])
            }
            
            # 创建pyrender场景
            scene = pyrender.Scene(ambient_light=[0.1, 0.1, 0.1])
            
            # 将trimesh网格转换为pyrender网格
            material = pyrender.MetallicRoughnessMaterial(
                metallicFactor=0.2,
                roughnessFactor=0.8,
                baseColorFactor=[0.8, 0.8, 0.8, 1.0]
            )
            
            # 处理可能的多网格情况
            if hasattr(mesh, 'vertices') and hasattr(mesh, 'faces'):
                # 单一网格情况
                mesh_node = pyrender.Mesh.from_trimesh(mesh, material=material)
                scene.add(mesh_node)
            else:
                # 尝试处理复杂网格
                try:
                    mesh_node = pyrender.Mesh.from_trimesh(mesh)
                    scene.add(mesh_node)
                except:
                    # 如果失败，使用备用处理方式
                    render_mesh = pyrender.Mesh.from_points(
                        mesh.vertices, 
                        colors=np.random.uniform(0, 1, size=(len(mesh.vertices), 3))
                    )
                    scene.add(render_mesh)
            
            # 添加光源
            light = DirectionalLight(color=[1.0, 1.0, 1.0], intensity=1.0)
            light_pose = np.eye(4)
            light_pose[:3, 3] = [0, 0, 10]
            scene.add(light, pose=light_pose)
            
            # 添加相机
            camera = PerspectiveCamera(yfov=np.pi/3.0)
            
            # 计算模型的边界框以确定相机位置
            bbox_max = mesh.bounds[1]
            bbox_min = mesh.bounds[0]
            scale = np.linalg.norm(bbox_max - bbox_min)
            center = (bbox_max + bbox_min) / 2
            
            # 对每个视角进行渲染
            for view_name, rotation in view_angles.items():
                try:
                    # 创建相机位姿 (从不同角度)
                    camera_pose = np.eye(4)
                    
                    # 计算相机位置 - 距离模型足够远以完全显示
                    camera_distance = scale * 2.5
                    camera_pose[:3, 3] = center + [
                        camera_distance * np.sin(rotation[0]) * np.cos(rotation[1]),
                        camera_distance * np.sin(rotation[0]) * np.sin(rotation[1]),
                        camera_distance * np.cos(rotation[0])
                    ]
                    
                    # 设置相机看向模型中心
                    camera_pose = self._look_at(camera_pose[:3, 3], center, np.array([0, 0, 1]))
                    
                    # 渲染图像
                    scene.add(camera, pose=camera_pose)
                    color, _ = self.renderer.render(scene)
                    
                    # 保存截图
                    screenshot_path = os.path.join(self.screenshot_dir, f"{view_name}.png")
                    plt.imsave(screenshot_path, color)
                    screenshots[view_name] = screenshot_path
                    
                    # 移除相机，准备下一个视角
                    scene.remove_node(list(scene.get_nodes())[-1])
                    
                except Exception as e:
                    logger.error(f"渲染{view_name}视角失败: {str(e)}")
                    # 如果某个视角渲染失败，使用模拟URL
                    screenshots[view_name] = f"{self.mock_screenshot_url}?angle={view_name}"
                    render_status = "partial_error"
            
            screenshots["render_status"] = render_status
            return screenshots
            
        except Exception as e:
            logger.error(f"多角度渲染失败: {str(e)}")
            return self._mock_render("", error=str(e))
    
    def _look_at(self, camera_position: np.ndarray, target: np.ndarray, up: np.ndarray) -> np.ndarray:
        """
        计算相机的look-at矩阵
        """
        # 计算相机方向向量
        forward = target - camera_position
        forward = forward / np.linalg.norm(forward)
        
        # 计算相机的right向量
        right = np.cross(forward, up)
        right = right / np.linalg.norm(right)
        
        # 重新计算up向量，确保三个向量正交
        up = np.cross(right, forward)
        
        # 构建旋转矩阵
        rotation = np.eye(4)
        rotation[0, :3] = right
        rotation[1, :3] = up
        rotation[2, :3] = -forward
        
        # 构建平移矩阵
        translation = np.eye(4)
        translation[:3, 3] = -camera_position
        
        # 组合得到look-at矩阵
        return rotation @ translation
    
    def _mock_render(self, model_url: str, error: str = None) -> Dict[str, Any]:
        """
        提供模拟的渲染结果（当无法使用实际渲染时）
        """
        mock_result = {
            "front": f"{self.mock_screenshot_url}?angle=front",
            "side": f"{self.mock_screenshot_url}?angle=side",
            "top": f"{self.mock_screenshot_url}?angle=top",
            "render_status": "success" if not error else "error"
        }
        
        if error:
            mock_result["error_message"] = error
            logger.warning(f"使用模拟渲染结果，错误: {error}")
        else:
            logger.warning("使用模拟渲染结果，因为pyrender不可用")
        
        return mock_result

# 豆包大模型客户端
class DoubaoClient:
    def __init__(self, api_key: str = "8e61af6d-bcbe-4b54-ac15-59761e7bd4cc"):
        self.api_key = api_key
        self.base_url = "https://api.doubao.com/api/v2"
        self.headers = {
            "Content-Type": "application/json",
            "Authorization": f"Bearer {self.api_key}"
        }
    
    def evaluate_model(self, prompt: str, image_urls: list, reference_image: Optional[str] = None) -> Dict[str, Any]:
        """
        调用豆包大模型评估3D模型
        
        Args:
            prompt: 用户提示词
            image_urls: 3D模型截图URL列表
            reference_image: 参考图片URL
            
        Returns:
            评估结果
        """
        try:
            # 构建评估提示词
            evaluation_prompt = self._build_evaluation_prompt(prompt, reference_image)
            
            # 准备请求数据
            data = {
                "model": "Doubao-Seed-1.6-vision",
                "messages": [
                    {
                        "role": "user",
                        "content": [
                            {"type": "text", "text": evaluation_prompt}
                        ]
                    }
                ],
                "temperature": 0.7,
                "max_tokens": 4000,
                "stream": True  # 使用SSE方式
            }
            
            # 添加图片数据
            for image_url in image_urls:
                data["messages"][0]["content"].append({
                    "type": "image",
                    "image_url": {"url": image_url}
                })
            
            logger.info("开始调用豆包大模型进行评估")
            
            # 发送请求（使用SSE）
            response = requests.post(
                f"{self.base_url}/chat/completions",
                headers=self.headers,
                json=data,
                stream=True
            )
            
            # 处理SSE响应
            full_response = ""
            for chunk in response.iter_lines():
                if chunk:
                    chunk = chunk.decode('utf-8')
                    if chunk.startswith('data: '):
                        chunk_data = chunk[6:]
                        if chunk_data != '[DONE]':
                            try:
                                json_data = json.loads(chunk_data)
                                if 'choices' in json_data and json_data['choices']:
                                    content = json_data['choices'][0].get('delta', {}).get('content', '')
                                    full_response += content
                            except json.JSONDecodeError:
                                continue
            
            logger.info("豆包大模型评估完成")
            
            # 解析评估报告
            return self._parse_evaluation_report(full_response)
            
        except Exception as e:
            logger.error(f"调用豆包大模型失败: {str(e)}")
            # 返回模拟评估结果
            return self._generate_mock_evaluation(prompt)
    
    def _build_evaluation_prompt(self, user_prompt: str, reference_image: Optional[str]) -> str:
        """构建评估提示词"""
        prompt = f"""
        请对以下3D模型进行专业评估。评估需要基于用户的原始提示词与生成的3D模型截图进行对比分析。
        
        用户原始提示词: {user_prompt}
        
        {f'参考图片: {reference_image}' if reference_image else ''}
        
        请从以下几个维度进行全面评估：
        1. 准确性：3D模型是否准确反映了用户提示词中的关键特征？
        2. 完整性：3D模型是否包含了所有必要的组成部分？
        3. 视觉质量：3D模型的纹理、颜色、形状是否自然合理？
        4. 比例与透视：3D模型的比例是否协调，透视是否正确？
        5. 创新与风格：3D模型是否有独特的艺术风格或创新点？
        
        请使用Markdown格式输出评估报告，包含：
        - 整体评分（1-10分）
        - 各维度详细分析
        - 优缺点总结
        - 改进建议
        
        请确保报告内容专业、详细且有建设性。
        """
        return prompt
    
    def _parse_evaluation_report(self, response: str) -> Dict[str, Any]:
        """解析评估报告"""
        # 实际应用中应该根据豆包返回的格式进行解析
        return {
            "title": "3D模型评估报告",
            "markdown_report": response,
            "overall_score": 8.5,  # 模拟评分
            "parsed_at": datetime.now().isoformat()
        }
    
    def _generate_mock_evaluation(self, prompt: str) -> Dict[str, Any]:
        """生成模拟评估结果"""
        mock_report = f"""
# 3D模型评估报告

## 整体评分
**8.5/10**

## 详细评估

### 1. 准确性 (8/10)
模型较好地反映了用户提示词中的主要特征，但在某些细节方面还有提升空间。

### 2. 完整性 (9/10)
模型包含了所有必要的组成部分，结构完整。

### 3. 视觉质量 (8/10)
纹理和颜色表现良好，但细节纹理可以更加丰富。

### 4. 比例与透视 (9/10)
模型比例协调，透视效果自然。

### 5. 创新与风格 (8/10)
具有一定的艺术风格，但可以更加突出独特性。

## 优缺点总结

### 优点
- 整体造型符合用户需求
- 比例协调，视觉效果良好
- 细节处理较为到位

### 缺点
- 部分细节纹理可以更加丰富
- 风格特色可以更加鲜明

## 改进建议
- 增加表面细节纹理
- 强化模型的独特风格特征
- 进一步优化光照效果
"""
        
        return {
            "title": "3D模型评估报告",
            "markdown_report": mock_report,
            "overall_score": 8.5,
            "generated_at": datetime.now().isoformat()
        }

# 3D模型评估主类
class ModelEvaluator:
    def __init__(self):
        self.renderer = ModelRenderer()
        self.doubao_client = DoubaoClient()
    
    def evaluate(self, model_url: str, prompt: str, user_id: int, reference_image: Optional[str] = None) -> Dict[str, Any]:
        """
        执行完整的3D模型评估流程
        
        Args:
            model_url: 3D模型URL
            prompt: 用户提示词
            user_id: 用户ID
            reference_image: 参考图片URL
            
        Returns:
            评估结果和报告
        """
        try:
            # 1. 渲染3D模型获取截图
            render_result = self.renderer.render_model(model_url)
            
            # 收集截图URL
            screenshot_urls = [
                render_result.get("front"),
                render_result.get("side"),
                render_result.get("top")
            ]
            screenshot_urls = [url for url in screenshot_urls if url]
            
            # 2. 调用豆包大模型进行评估
            evaluation_result = self.doubao_client.evaluate_model(
                prompt=prompt,
                image_urls=screenshot_urls,
                reference_image=reference_image
            )
            
            # 3. 整合结果
            result = {
                "title": evaluation_result["title"],
                "markdown_report": evaluation_result["markdown_report"],
                "model_3d_url": model_url,
                "user_id": user_id,
                "screenshot": screenshot_urls[0] if screenshot_urls else None,
                "overall_score": evaluation_result.get("overall_score"),
                "render_status": render_result.get("render_status", "unknown"),
                "evaluation_status": "completed"
            }
            
            logger.info(f"3D模型评估流程完成，用户ID: {user_id}")
            return result
            
        except Exception as e:
            logger.error(f"3D模型评估失败: {str(e)}")
            # 发生异常时返回基本信息
            return {
                "title": "评估失败报告",
                "markdown_report": f"# 评估失败\n\n由于系统原因，3D模型评估未能完成。\n\n错误信息: {str(e)}",
                "model_3d_url": model_url,
                "user_id": user_id,
                "screenshot": None,
                "evaluation_status": "failed",
                "error_message": str(e)
            }
    
    def evaluate_async(self, model_url: str, prompt: str, user_id: int, callback=None, reference_image: Optional[str] = None):
        """
        异步执行3D模型评估
        
        Args:
            model_url: 3D模型URL
            prompt: 用户提示词
            user_id: 用户ID
            callback: 评估完成后的回调函数
            reference_image: 参考图片URL
        """
        def _evaluate_task():
            result = self.evaluate(model_url, prompt, user_id, reference_image)
            if callback:
                callback(result)
            return result
        
        # 使用线程池执行异步评估
        with ThreadPoolExecutor(max_workers=1) as executor:
            future = executor.submit(_evaluate_task)
            return future